package userclient

import (
	"io"
	"log"
	"net"
	"strconv"
	"sync"

	"gitee.com/klen-ygs/portmap/udpblock"
)

type UdpBlock struct {
	data []byte
	addr *net.UDPAddr
}

type Instance struct {
	config         Config
	errChannel     chan error
	udpSendChannel chan UdpBlock
	udpSendMap     sync.Map
	tcpListener    net.Listener
	udpListener    *net.UDPConn
}

// NewInstance 返回一个用户流量转发实例
func NewInstance(config Config) *Instance {
	return &Instance{
		config:         config,
		errChannel:     make(chan error, 10),
		udpSendChannel: make(chan UdpBlock, 20),
	}
}

// Run 将运行该用户转发实例，当该示例正常运行时，会一直阻塞，直到发生错误
func (i *Instance) Run() {
	wg := &sync.WaitGroup{}
	wg.Add(2)

	go i.listenTcp(wg)
	go i.listenUdp(wg)
	i.waitClose(wg)

}

// waitClose 将一直监听 i.errChannel 是否有错误信息，若该实例发生错误，则清理资源
func (i *Instance) waitClose(wg *sync.WaitGroup) {
	wg.Wait()
	err := <-i.errChannel
	log.Printf("[Error] 映射实例被关闭:%s", err.Error())
	i.tcpListener.Close()
	i.udpListener.Close()
}

// tcpToTcp 将 src 数据流向另一个数据流 dst 转发，若发生错误，则向 enChan 发送错误
func (i *Instance) tcpToTcp(endChan chan error, src io.Reader, dst io.Writer) {
	_, err := io.Copy(dst, src)
	endChan <- err
}

// srvTcp 获得到用户连接 conn 后，将建立一个新tcp连接到服务器，地址是 i.config.ServerTcpEndpoint，
// 连接建立之后，将流量在两个数据流中转发
func (i *Instance) srvTcp(conn net.Conn) {
	defer conn.Close()

	dial, err := net.Dial("tcp", i.config.ServerTcpEndpoint)
	if err != nil {
		log.Printf("[Error] 连接服务器失败:%s", err.Error())
		return
	}
	defer dial.Close()

	endChan := make(chan error, 2)
	go i.tcpToTcp(endChan, conn, dial)
	go i.tcpToTcp(endChan, dial, conn)

	err = <-endChan
	log.Printf("[info] TCP连接%s->%s关闭:%s", conn.RemoteAddr(), dial.RemoteAddr(), err.Error())
}

// listenTcp 将一直监听用户的tcp连接，监听端口是 i.config.ListenPort
func (i *Instance) listenTcp(wg *sync.WaitGroup) {
	listen, err := net.Listen("tcp", ":"+strconv.Itoa(i.config.ListenPort))
	if err != nil {
		wg.Done()
		log.Printf("[Error] 监听tcp连接失败")
		return
	}
	i.tcpListener = listen
	wg.Done()

	for {
		accept, err := listen.Accept()
		if err != nil {
			i.errChannel <- err
			return
		}
		log.Printf("[info] 新连接接入:%s", accept.RemoteAddr())
		go i.srvTcp(accept)
	}

}

// getUdpDataChannel 会从 udpSendMap 中以 addr 为key值找到发送数据的channel，如果没有该channel返回false
func (i *Instance) getUdpDataChannel(addr string) (chan []byte, bool) {
	value, ok := i.udpSendMap.Load(addr)
	if !ok {
		return nil, false
	}
	return value.(chan []byte), ok
}

// getDataSendUdpToLocal 会从 i.udpSendChannel 中读取要发送的信息，并发送到给用户
func (i *Instance) getDataSendUdpToLocal() {
	for {
		send := <-i.udpSendChannel
		_, err := i.udpListener.WriteToUDP(send.data, send.addr)
		if err != nil {
			i.errChannel <- err
			return
		}
	}
}

// readAndSendLocalUdp 会从 tcp 中读取数据并转发到 udpSendChannel
func (i *Instance) readAndSendLocalUdp(endChan chan error, tcp io.Reader, addr *net.UDPAddr) {
	buf := make([]byte, 2<<16)
	for {
		n, err := tcp.Read(buf)
		if err != nil {
			endChan <- err
			break
		}

		send := make([]byte, n)
		copy(send, buf[:n])
		i.udpSendChannel <- UdpBlock{
			data: send,
			addr: addr,
		}
	}

}

// sendDataToServer 会从 dataChan 中读取本机发送的udp信息，并将数据转发到服务器
func (i *Instance) sendDataToServer(endChan chan error, tcp io.Writer, dataChan chan []byte) {
	for {
		send, ok := <-dataChan
		if !ok {
			return
		}
		_, err := tcp.Write(send)
		if err != nil {
			endChan <- err
			break
		}
	}
}

// newUdpSender 会新建一个udp转发连接，去向服务器的路线是 dataChan -> server
// 去向本机的路线是 server -> udpSendChannel
func (i *Instance) newUdpSender(addr *net.UDPAddr, dataChan chan []byte) {
	dial, err := net.Dial("tcp", i.config.ServerUdpEndpoint)
	if err != nil {
		log.Printf("[Error] 连接服务器udp失败")
		return
	}
	defer dial.Close()

	writer, err := udpblock.NewDataWriter(dial)
	if err != nil {
		log.Printf("[Error] 构建tcp块写入对象失败")
		return
	}

	reader, err := udpblock.NewDataReader(dial)
	if err != nil {
		log.Printf("[Error] 构建tcp块读取对象失败")
		return
	}
	defer reader.Close()

	endChan := make(chan error, 2)
	go i.readAndSendLocalUdp(endChan, reader, addr)
	go i.sendDataToServer(endChan, writer, dataChan)

	err = <-endChan
	// 关闭传输通道并删除该地址映射
	channel, find := i.getUdpDataChannel(addr.String())
	if find {
		close(channel)
	}
	i.udpSendMap.Delete(addr.String())
	log.Printf("[info] UDP连接：%s->%s关闭:%s", addr.String(), dial.RemoteAddr(), err.Error())
}

// listenUdp 将不断监听本地发送的udp包，若之前已经从相同地址发送过udp包，则在 getUdpDataChannel
// 获取发送数据的channel来发送数据，否则将创建一个新的channel，并发起新的连接到服务器
func (i *Instance) listenUdp(wg *sync.WaitGroup) {
	udp, err := net.ListenUDP("udp", &net.UDPAddr{
		Port: i.config.ListenPort,
	})
	if err != nil {
		wg.Done()
		i.errChannel <- err
		return
	}
	i.udpListener = udp
	wg.Done()
	go i.getDataSendUdpToLocal()

	buf := make([]byte, 1500)
	for {
		n, addr, err := udp.ReadFromUDP(buf)
		if err != nil {
			i.errChannel <- err
			return
		}

		if channel, ok := i.getUdpDataChannel(addr.String()); ok {
			send := make([]byte, n)
			copy(send, buf[:n])
			channel <- send
		} else {
			log.Printf("[info] udp新连接接入:%s", addr.String())
			send := make([]byte, n)
			copy(send, buf[:n])
			sendChan := make(chan []byte, 10)
			sendChan <- send
			i.udpSendMap.Store(addr.String(), sendChan)
			go i.newUdpSender(addr, sendChan)
		}

	}

}
